#include <iostream>
#include <stdio.h>
#include <opencv.hpp>
#include <highgui/highgui.hpp>
#include <rawfile.h>

using namespace std;

const int BOARDW = 11;
const int BOARDH = 8;
const int IMGWIDTH = 640;
const int IMGHEIGHT = 480;
const int TOTALCAM = 24;
const int BLOCKSIZE = IMGWIDTH * IMGHEIGHT * TOTALCAM;
const int MAXFRAME = 75;
const int BOARDBLOCKWIDTH = 80;
const int BOARDBLOCKHEIGHT = 80;
const float PIXELWIDTH = 0.006f;
const float PIXELHEIGHT = 0.006f;

//rawfile path, output path, camera number

int main(int argc, char* argv[])
{
    if(argc != 4)
    {
        cout << "raw file path, output path, camera number" << endl;
        return 0;
    }
    int board_n = BOARDW*BOARDH;
    CvSize board_sz = cvSize(BOARDW, BOARDH);

    char rawPath[1024];
    snprintf(rawPath, sizeof(rawPath),"%s", argv[1]);
    PFCMU::RAWFile raw;
    raw.open(rawPath,IMGWIDTH,IMGHEIGHT);
    off64_t frameCount = PFCMU::RAWFile::framecount(rawPath, BLOCKSIZE);

    char desPath[1024];
    snprintf(desPath,sizeof(desPath),"%s",argv[2]);
    char *prefix = "img";

    int cami = atoi(argv[3]);
    cami = cami-1;

    IplImage *rawImage = cvCreateImage(cvSize(IMGWIDTH,IMGHEIGHT),IPL_DEPTH_8U,1);
    IplImage *rgbImage = cvCreateImage(cvSize(IMGWIDTH,IMGHEIGHT),IPL_DEPTH_8U,3);
    IplImage *grayImage = cvCreateImage(cvSize(IMGWIDTH,IMGHEIGHT),IPL_DEPTH_8U,1);


    //frameCount = 1000;
    //for(int cami=0;cami<TOTALCAM;cami++)
    //{
        cout << "Start camera " << cami+1 << "/" << TOTALCAM << endl;
        char camName[128];
        snprintf(camName,sizeof(camName),"cam%d",cami+1);
        //every cam
        vector<CvPoint2D32f *> cornervec;
        int successes = 0;
        for(int framei=0;framei<frameCount;framei+=3)
        {
            cout << "Processing camera " << cami+1 << " frame " << framei+1 << "/" << frameCount << endl;
            //every frame
            cvZero(rawImage);cvZero(rgbImage);cvZero(grayImage);
            PFCMU::RAWFile::const_iterator itr = raw.at(framei);
            itr.extract(cami,rawImage);
            cvCvtColor(rawImage,rgbImage,CV_BayerGR2BGR);

            CvPoint2D32f *corners = new CvPoint2D32f[board_n];
            int corner_count = 0;
            int rawFound = cvFindChessboardCorners(rgbImage,board_sz,corners,&corner_count,
                                                   CV_CALIB_CB_ADAPTIVE_THRESH|CV_CALIB_CB_FILTER_QUADS);
            if(rawFound==0)
            {
                delete[] corners;
                corners = NULL;
                continue;
            }
            cvCvtColor(rgbImage,grayImage,CV_BGR2GRAY);
            cvFindCornerSubPix(grayImage,corners,corner_count,cvSize(5,5),cvSize(-1,-1),
                               cvTermCriteria(CV_TERMCRIT_EPS+CV_TERMCRIT_ITER,30,0.1));
            cvDrawChessboardCorners(rgbImage, board_sz, corners, corner_count, rawFound);
            if(corner_count == board_n)
            {
//                cvShowImage("SingleImage", rgbImage);
//                int c = cvWaitKey(0);
//                if(c=='a')
//                {
//                    //accept
//                }
//                else if(c=='x')
//                {
//                    cout << "User stopped the collection, continue to calibration" << endl;
//                    delete[] corners;
//                    break;
//                }
//                else
//                {
//                    framei+=5;
//                    continue;
//                }
                cornervec.push_back(corners);
                char desFinal[1024];
                snprintf(desFinal,sizeof(desFinal),"%s/%s/%s%05d.bmp", desPath,camName,prefix,framei);
                cvSaveImage(desFinal,rgbImage);
                successes++;
                framei+=10;
                if(MAXFRAME!=0 && successes>MAXFRAME)
                {
                    break;
                }
            }
            else
            {
                delete[] corners;
                corners = NULL;
                continue;
            }
            cout << "Found " << successes << " frames" << endl;
        }//for(int framei=0;framei<frameCount;framei++)
        cout << "There are totally " << successes << " frames on camera " << cami+1 << endl;
        if(successes==0)
        {
            cout << "Found 0 chessboard. Exit" << endl;
            cvReleaseImage(&rawImage);
            cvReleaseImage(&rgbImage);
            cvReleaseImage(&grayImage);
            return 0;
        }
        cout << "Copying data" << endl;
        CvMat *image_points = cvCreateMat(successes*board_n,2,CV_32FC1);
        CvMat *object_points = cvCreateMat(successes*board_n,3,CV_32FC1);
        CvMat *point_counts = cvCreateMat(successes,1,CV_32SC1);
        CvMat *intrinsic_matrix = cvCreateMat(3,3,CV_32FC1);
        CvMat *distortion_coeffs = cvCreateMat(5,1,CV_32FC1);
        for(int imgi=0;imgi<successes;imgi++)
        {
            int bitoffset = imgi*board_n;
            CvPoint2D32f *corners = cornervec.at(imgi);
            for(int pointi=0;pointi<board_n;pointi++)
            {
                CV_MAT_ELEM(*image_points, float, bitoffset+pointi,0) = corners[pointi].x;
                CV_MAT_ELEM(*image_points, float, bitoffset+pointi,1) = corners[pointi].y;
                CV_MAT_ELEM(*object_points, float, bitoffset+pointi,0) = pointi/BOARDW*BOARDBLOCKWIDTH;
                CV_MAT_ELEM(*object_points, float, bitoffset+pointi,1) = pointi%BOARDW*BOARDBLOCKHEIGHT;
                CV_MAT_ELEM(*object_points, float, bitoffset+pointi,2) = 0.0f;
            }//for(int pointi=0;pointi<board_n;pointi++)
            CV_MAT_ELEM(*point_counts, int, imgi, 0) = board_n;
        }//for(int imgi=0;imgi<successes;imgi++)

        CV_MAT_ELEM(*intrinsic_matrix, float, 0,0) = 1.0f;
        CV_MAT_ELEM(*intrinsic_matrix, float, 1,1) = 1.0f;
        cout << "Start calibration" << endl;
        cvCalibrateCamera2(object_points, image_points, point_counts,
                           cvGetSize(rgbImage), intrinsic_matrix,
                           distortion_coeffs, NULL, NULL, 0);
        cout << "Calibration finished! Save data" << endl;
        char intrinsicsPath[1024];
        char distortionPath[1024];
        snprintf(intrinsicsPath,sizeof(intrinsicsPath),"%s/%s/intrinsics.xml", desPath,camName);
        snprintf(distortionPath,sizeof(distortionPath),"%s/%s/distortion.xml", desPath,camName);
        cvSave(intrinsicsPath,intrinsic_matrix);
        cvSave(distortionPath, distortion_coeffs);
        for(int imgi=0;imgi<successes;imgi++)
        {
            CvPoint2D32f *corners = cornervec.at(imgi);
            delete[] corners;
        }
        cornervec.clear();
        cout << "Finished camera " << cami+1 << "/" << TOTALCAM << endl;
        cvReleaseMat(&object_points);
        cvReleaseMat(&image_points);
        cvReleaseMat(&point_counts);
        cvReleaseMat(&intrinsic_matrix);
        cvReleaseMat(&distortion_coeffs);
    //}//for(int cami=0;cami<TOTALCAM;cami++)

    cout << "Finished!" << endl;
    cvReleaseImage(&rawImage);
    cvReleaseImage(&rgbImage);
    cvReleaseImage(&grayImage);
    return 0;
}

